

import os
import random
import logging

import torch
import numpy as np
from sklearn.metrics import f1_score

from transformers import BertConfig
from transformers import BertTokenizer

from bert_finetune_cls.model import ClsBERT
# 模型字典映射
MODEL_CLASSES = {
    'bert': (BertConfig, ClsBERT, BertTokenizer),
}
# 模型路径字典映射
MODEL_PATH_MAP = {
    # 'bert': './bert_finetune_cls1/resources/bert_base_uncased',
    'bert': './bert_finetune_cls/resources/uncased_L-2_H-128_A-2',
}

# 获得意图标签
def get_intent_labels(args):
    return [label.strip() for label in open(os.path.join(args.data_dir, args.task, args.intent_label_file), 'r', encoding='utf-8')]

# 加载分词模型
def load_tokenizer(args):
    return MODEL_CLASSES[args.model_type][2].from_pretrained(args.model_name_or_path)

# 日志初始化
def init_logger():
    logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
                        datefmt='%m/%d/%Y %H:%M:%S',
                        level=logging.INFO)

# 固定种子
def set_seed(args):
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    if not args.no_cuda and torch.cuda.is_available():
        torch.cuda.manual_seed_all(args.seed)

# 打分
def compute_metrics(intent_preds, intent_labels):
    assert len(intent_preds) == len(intent_labels)
    results = {}
    intent_result = get_intent_acc(intent_preds, intent_labels)

    results.update(intent_result)

    return results

#
def get_intent_acc(preds, labels):
    acc = (preds == labels).mean()

    f1 = f1_score(
        labels, preds
    )

    return {
        "acc": acc,
        "f1_score": f1,
        "score": (acc + f1) / 2,
    }


# 获得预测文本
def read_prediction_text(args):
    return [text.strip() for text in open(os.path.join(args.pred_dir, args.pred_input_file), 'r', encoding='utf-8')]
